import pygame
import sys
import os
import datetime
import time
from ai_player import AIPlayer

# 导入MCTS+神经网络AI适配器
from ai_mcts_adapter import AIMCTSPlayer

# 游戏常量
BOARD_SIZE = 15  # 15x15棋盘
CELL_SIZE = 40  # 每个格子的大小
MARGIN = 40  # 边缘距离
SCREEN_WIDTH = CELL_SIZE * (BOARD_SIZE - 1) + MARGIN * 2
SCREEN_HEIGHT = SCREEN_WIDTH + 160  # 额外空间用于显示信息和按钮
# 渲染优化常量
LAST_MOVE_MARKER_RADIUS = 5  # 最后落子标记半径
LAST_MOVE_MARKER_COLOR = (255, 0, 0)  # 红色标记最后落子位置

# 颜色定义
BACKGROUND_COLOR = (240, 190, 100)  # 棋盘底色(棕黄色)
LINE_COLOR = (0, 0, 0)  # 线条颜色(黑色)
BLACK_STONE = (0, 0, 0)  # 黑棋
WHITE_STONE = (255, 255, 255)  # 白棋
TEXT_COLOR = (0, 0, 0)  # 文本颜色
BUTTON_COLOR = (180, 180, 180)  # 按钮颜色
BUTTON_HOVER = (150, 150, 150)  # 按钮悬停颜色
BUTTON_TEXT_COLOR = (0, 0, 0)  # 按钮文字颜色

# 游戏模式常量
MODE_HUMAN_VS_HUMAN = "human_vs_human"
MODE_HUMAN_VS_AI = "human_vs_ai"


class GameCore:
    """核心游戏模块：管理棋盘状态、胜负判定、游戏流程"""

    def __init__(self):
        self.reset_game()
        self.game_mode = MODE_HUMAN_VS_HUMAN  # 默认为双人对战
        self.ai_player = None
        self.ai_type = "traditional"  # 默认使用传统AI
        self.ai_thinking = False
        self.move_history = []  # 存储落子历史，用于悔棋
        self.ai_difficulty = "medium"  # 默认中等难度

    def reset_game(self):
        """重置游戏至初始状态"""
        self.board = [[0 for _ in range(BOARD_SIZE)] for _ in range(BOARD_SIZE)]  # 15×15棋盘，0:空，1:黑棋，2:白棋
        self.current_player = 1  # 当前玩家（1:黑，2:白）
        self.game_over = False  # 游戏是否结束
        self.winner = 0  # 获胜方（0:无，1:黑，2:白）
        self.message = "黑方回合"  # 状态提示信息
        self.move_history = []  # 清空历史记录
        self.placed_stones = []  # 新增：记录已下棋子的位置和玩家
        self.stone_count = 0  # 新增：记录已下棋子数量，用于快速平局检查
        self.last_move = None  # 新增：记录上次落子位置
        self.ai_thinking = False

    def set_game_mode(self, mode, ai_difficulty=None):
        """设置游戏模式和AI难度
        
        Args:
            mode: 游戏模式，MODE_HUMAN_VS_HUMAN 或 MODE_HUMAN_VS_AI
            ai_difficulty: AI难度级别，"easy", "normal", "medium", "hard", "expert"
        """
        self.game_mode = mode
        self.reset_game()
        
        if mode == MODE_HUMAN_VS_AI:
            # 支持5个难度级别，默认使用medium
            self.ai_difficulty = ai_difficulty or "medium"
            
            # 询问用户选择AI类型
            import tkinter as tk
            from tkinter import messagebox
            root = tk.Tk()
            root.withdraw()  # 隐藏主窗口
            use_mcts = messagebox.askyesno(
                "选择AI类型", 
                "是否使用MCTS+神经网络AI？\n(是 = MCTS+神经网络AI, 否 = 传统AI)",
                default="yes"
            )
            root.destroy()
            
            # 创建AI玩家
            if use_mcts:
                self.ai_type = "mcts"
                # 使用MCTS+神经网络AI适配器
                self.ai_player = AIMCTSPlayer(difficulty=self.ai_difficulty)
                self.ai_player.player = 2  # AI执白
                self.ai_player.human_player = 1
                messagebox.showinfo("AI类型", "已选择MCTS+神经网络AI")
            else:
                self.ai_type = "traditional"
                # 使用传统AI
                self.ai_player = AIPlayer(difficulty=self.ai_difficulty)
                self.ai_player.player = 2  # AI执白
                self.ai_player.human_player = 1
            
            self.message = "黑方回合（您）"

    def make_move(self, row, col):
        """验证落子有效性并更新棋盘状态"""
        if self.game_over:
            self.message = "游戏已结束，请重新开始"
            return False

        # 检查位置是否为空
        if self.board[row][col] != 0:
            self.message = "此处已有棋子，请重新选择"
            return False

        # 记录落子历史
        self.move_history.append((row, col, self.current_player))
        
        # 新增：记录已放置的棋子
        self.placed_stones.append((row, col, self.current_player))
        self.stone_count += 1
        
        # 记录最后落子位置
        self.last_move = (row, col)
        
        # 落子
        self.board[row][col] = self.current_player

        # 检查是否获胜
        if self.check_win(row, col):
            self.game_over = True
            self.winner = self.current_player
            self.message = f"{'黑方' if self.winner == 1 else '白方'}获胜！"
            return True

        # 优化：使用计数器进行平局检查，O(1)复杂度
        if self.stone_count >= BOARD_SIZE * BOARD_SIZE:
            self.game_over = True
            self.message = "平局！"
            return True

        # 切换玩家
        self.current_player = 2 if self.current_player == 1 else 1
        
        # 更新消息
        if self.game_mode == MODE_HUMAN_VS_AI:
            if self.current_player == 1:
                self.message = "黑方回合（您）"
            else:
                self.message = "AI思考中..."
        else:
            self.message = f"{'黑方' if self.current_player == 1 else '白方'}回合"
        
        return True

    def check_win(self, row, col):
        """优化版胜负判定算法，减少边界检查次数"""
        directions = [
            (0, 1),  # 水平向右
            (1, 0),  # 垂直向下
            (1, 1),  # 右下对角线
            (1, -1)  # 左下对角线
        ]

        player = self.board[row][col]

        for dr, dc in directions:
            count = 1  # 已包含当前落子
            
            # 正方向检查
            r, c = row + dr, col + dc
            while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and self.board[r][c] == player:
                count += 1
                if count >= 5:  # 提前返回，发现五连珠立即返回
                    return True
                r += dr
                c += dc
            
            # 反方向检查
            r, c = row - dr, col - dc
            while 0 <= r < BOARD_SIZE and 0 <= c < BOARD_SIZE and self.board[r][c] == player:
                count += 1
                if count >= 5:  # 提前返回
                    return True
                r -= dr
                c -= dc
            
            if count >= 5:
                return True
        return False
        
    def can_undo(self):
        """检查是否可以悔棋
        
        Returns:
            bool: 是否有可撤销的步骤
        """
        return len(self.move_history) > 0
        
    def ai_move(self):
        """优化版AI自动落子算法（多线程实现）
        
        Returns:
            bool: 是否成功落子
        """
        if self.game_over or self.game_mode != MODE_HUMAN_VS_AI or self.current_player != self.ai_player.player:
            return False
            
        self.ai_thinking = True
        
        # 创建线程执行AI计算
        import threading
        ai_thread = threading.Thread(target=self._ai_think)
        ai_thread.daemon = True  # 设为守护线程，避免阻止程序退出
        ai_thread.start()
        
        return True
        
    def _ai_think(self):
        """优化版AI思考逻辑，加入难度级别控制和搜索优化
        
        使用pygame事件系统通知主线程落子结果
        """
        try:
            # 根据难度级别设置思考时间和搜索参数
            difficulty_settings = {
                "easy": {"max_time": 0.5, "search_depth": 1, "use_alpha_beta": False, "prune_factor": 1.0, "max_candidates": 8},
                "normal": {"max_time": 1.0, "search_depth": 2, "use_alpha_beta": True, "prune_factor": 0.9, "max_candidates": 12},
                "medium": {"max_time": 1.5, "search_depth": 3, "use_alpha_beta": True, "prune_factor": 0.8, "max_candidates": 15},
                "hard": {"max_time": 2.5, "search_depth": 4, "use_alpha_beta": True, "prune_factor": 0.7, "max_candidates": 20},
                "expert": {"max_time": 4.0, "search_depth": 5, "use_alpha_beta": True, "prune_factor": 0.6, "max_candidates": 25}
            }
            
            # 获取当前难度设置
            settings = difficulty_settings.get(self.ai_difficulty, difficulty_settings["medium"])
            
            # 根据AI类型调用不同的查找最佳落子方法
            if self.ai_type == "mcts":
                # MCTS+神经网络AI
                best_move = self.ai_player.find_best_move(
                    self.board,
                    max_time=settings["max_time"]
                )
            else:
                # 传统AI
                # 设置AI搜索深度
                self.ai_player.set_search_depth(settings["search_depth"])
                
                # 执行AI搜索，传入所有难度相关参数
                best_move = self.ai_player.find_best_move(
                    self.board, 
                    max_time=settings["max_time"],
                    use_alpha_beta=settings["use_alpha_beta"],
                    prune_factor=settings["prune_factor"],
                    max_candidates=settings["max_candidates"]
                )
            
            if best_move and not self.game_over:  # 确保游戏未结束
                # 回到主线程执行落子
                import pygame
                pygame.event.post(pygame.event.Event(pygame.USEREVENT, move=best_move))
        except Exception as e:
            print(f"AI思考过程中出错: {e}")
        finally:
            self.ai_thinking = False
    
    def undo_move(self):
        """优化版悔棋功能，撤销上一步操作
        
        Returns:
            bool: 是否成功悔棋
        """
        if not self.can_undo():
            self.message = "没有可悔的步骤"
            return False
            
        # 撤销两步（玩家和AI各一步）如果是人机对战且当前是AI回合
        steps_to_undo = 1
        if self.game_mode == MODE_HUMAN_VS_AI:
            # 如果最后一步是AI下的，需要撤销两步
            if len(self.move_history) >= 2 and self.move_history[-1][2] == self.ai_player.player:
                steps_to_undo = 2
            else:
                steps_to_undo = 1
        
        # 执行撤销
        for _ in range(min(steps_to_undo, len(self.move_history))):
            row, col, _ = self.move_history.pop()
            self.board[row][col] = 0
            # 维护placed_stones列表
            if self.placed_stones:  # 确保列表不为空
                self.placed_stones.pop()
            self.stone_count = max(0, self.stone_count - 1)  # 确保计数不为负
        
        # 重置游戏状态
        self.game_over = False
        self.winner = 0
        
        # 更新最后落子位置
        self.last_move = self.move_history[-1][:2] if self.move_history else None
        
        # 更新当前玩家
        if self.move_history:
            # 当前玩家为最后一步的下一个玩家
            last_player = self.move_history[-1][2]
            self.current_player = 2 if last_player == 1 else 1
        else:
            self.current_player = 1
        
        # 更新消息
        if self.game_mode == MODE_HUMAN_VS_AI:
            if self.current_player == 1:
                self.message = "黑方回合（您）"
            else:
                self.message = "AI思考中..."
        else:
            self.message = f"{'黑方' if self.current_player == 1 else '白方'}回合"
        
        return True
    



# AIPlayer类已从ai_player模块导入


class InteractionHandler:
    """交互处理模块：处理鼠标/键盘输入"""

    def __init__(self, game_core):
        self.game_core = game_core  # 关联核心游戏模块
        self.keyboard_input = ""  # 存储键盘输入字符串
        self.waiting_for_input = False  # 是否处于键盘输入模式
        self.waiting_for_filename = False  # 是否处于等待用户输入文件名的状态
        self.filename_input = ""  # 存储用户输入的文件名
        self.data_manager = None  # 数据管理模块引用
        self.waiting_for_load_filename = False  # 是否等待加载文件名输入

    def set_data_manager(self, data_manager):
        """设置数据管理模块引用"""
        self.data_manager = data_manager

    def handle_mouse_click(self, pos):
        """将鼠标坐标转换为棋盘索引并触发落子"""
        row, col = self.convert_screen_to_board(pos)
        if row is not None and col is not None:
            self.game_core.make_move(row, col)
        else:
            self.game_core.message = "请点击棋盘交叉点"

    def convert_screen_to_board(self, pos):
        """将屏幕像素坐标转换为棋盘行列索引"""
        x, y = pos

        # 计算理论行列索引
        row = round((y - MARGIN) / CELL_SIZE)
        col = round((x - MARGIN) / CELL_SIZE)

        # 检查是否在有效范围内
        if 0 <= row < BOARD_SIZE and 0 <= col < BOARD_SIZE:
            # 检查是否离交叉点足够近（容错范围15像素）
            actual_x = MARGIN + col * CELL_SIZE
            actual_y = MARGIN + row * CELL_SIZE
            if abs(x - actual_x) < 15 and abs(y - actual_y) < 15:
                return row, col

        return None, None

    def handle_keyboard_event(self, event):
        """处理键盘输入事件"""
        if self.waiting_for_filename:
            self.handle_filename_input(event)
        elif self.waiting_for_load_filename:
            self.handle_load_filename_input(event)
        elif event.key == pygame.K_RETURN:  # 回车键
            if self.waiting_for_input:
                # 提交输入内容并处理
                self.parse_keyboard_input()
                self.keyboard_input = ""
                self.waiting_for_input = False
            else:
                # 进入输入模式
                self.waiting_for_input = True
                self.game_core.message = "请输入落子位置(行 列)，按回车确认（例如：3 5）"
        elif event.key == pygame.K_BACKSPACE:  # 退格键
            if self.waiting_for_input:
                self.keyboard_input = self.keyboard_input[:-1]
        elif event.key == pygame.K_ESCAPE:  # ESC键，退出输入模式
            if self.waiting_for_input:
                self.keyboard_input = ""
                self.waiting_for_input = False
                self.game_core.message = f"{'黑方' if self.game_core.current_player == 1 else '白方'}回合"
        elif (pygame.K_0 <= event.key <= pygame.K_9 or event.key == pygame.K_SPACE) and self.waiting_for_input:
            # 只允许输入数字和空格
            self.keyboard_input += event.unicode

    def parse_keyboard_input(self):
        """解析键盘输入的行号列号并触发落子"""
        if not self.keyboard_input.strip():
            self.game_core.message = "输入不能为空，请重新输入"
            return

        # 分割输入字符串
        parts = self.keyboard_input.strip().split()
        if len(parts) != 2:
            self.game_core.message = "格式错误，请输入'行号 列号'（如3 5）"
            return

        try:
            # 转换为整数（玩家输入从1开始，转换为0开始的索引）
            row = int(parts[0]) - 1
            col = int(parts[1]) - 1

            # 检查是否在有效范围内
            if not (0 <= row < BOARD_SIZE and 0 <= col < BOARD_SIZE):
                self.game_core.message = f"超出范围，请输入1-{BOARD_SIZE}之间的行/列号"
                return

            # 尝试落子
            self.game_core.make_move(row, col)

        except ValueError:
            self.game_core.message = "请输入数字"

    def handle_filename_input(self, event):
        """处理用户输入的文件名"""
        if event.key == pygame.K_RETURN:  # 确认输入
            if self.filename_input.strip():
                # 调用数据管理模块保存棋局
                try:
                    saved_filename = self.data_manager.save_game(
                        self.game_core,
                        self.filename_input.strip()
                    )
                    self.game_core.message = f"棋局已保存: {saved_filename}"
                except Exception as e:
                    self.game_core.message = str(e)
            else:
                # 无输入时使用自动生成的文件名
                try:
                    saved_filename = self.data_manager.save_game(self.game_core)
                    self.game_core.message = f"棋局已保存: {saved_filename}"
                except Exception as e:
                    self.game_core.message = str(e)
            self.waiting_for_filename = False
            self.filename_input = ""
        elif event.key == pygame.K_BACKSPACE:  # 删除字符
            self.filename_input = self.filename_input[:-1]
        elif event.key == pygame.K_ESCAPE:  # 取消输入
            self.waiting_for_filename = False
            self.filename_input = ""
            self.game_core.message = "取消保存"
        else:  # 输入字符
            self.filename_input += event.unicode

    def handle_load_filename_input(self, event):
        """处理用户输入的加载文件名"""
        if event.key == pygame.K_RETURN:  # 确认输入
            if self.filename_input.strip():
                # 调用数据管理模块加载棋局
                try:
                    result = self.data_manager.load_game_by_filename(self.filename_input.strip())
                    if result:
                        self.game_core.board = result['board']
                        self.game_core.current_player = result['current_player']
                        self.game_core.game_over = result['game_over']
                        self.game_core.winner = result['winner']
                        
                        # 根据棋盘状态重建其他属性
                        self.game_core.move_history = []
                        self.game_core.placed_stones = []
                        self.game_core.stone_count = 0
                        self.game_core.last_move = None
                        
                        # 遍历棋盘重建历史记录和其他属性
                        for row_idx, row in enumerate(self.game_core.board):
                            for col_idx, cell in enumerate(row):
                                if cell != 0:  # 非空位置
                                    self.game_core.move_history.append((row_idx, col_idx, cell))
                                    self.game_core.placed_stones.append((row_idx, col_idx, cell))
                                    self.game_core.stone_count += 1
                                    # 如果是最后一个落子，更新last_move
                                    if self.game_core.stone_count == len(self.game_core.move_history):
                                        self.game_core.last_move = (row_idx, col_idx)

                        if self.game_core.game_over:
                            if self.game_core.winner != 0:
                                self.game_core.message = f"已加载棋局: {result['filename']}，{'黑方' if self.game_core.winner == 1 else '白方'}获胜"
                            else:
                                self.game_core.message = f"已加载棋局: {result['filename']}，平局"
                        else:
                            self.game_core.message = f"已加载棋局: {result['filename']}，当前为{'黑方' if self.game_core.current_player == 1 else '白方'}回合"
                    else:
                        self.game_core.message = "取消加载"
                except Exception as e:
                    self.game_core.message = str(e)
            else:
                self.game_core.message = "文件名不能为空"
            self.waiting_for_load_filename = False
            self.filename_input = ""
        elif event.key == pygame.K_BACKSPACE:  # 删除字符
            self.filename_input = self.filename_input[:-1]
        elif event.key == pygame.K_ESCAPE:  # 取消输入
            self.waiting_for_load_filename = False
            self.filename_input = ""
            self.game_core.message = "取消加载"
        else:  # 输入字符
            self.filename_input += event.unicode


class UIRenderer:
    """界面渲染模块：绘制棋盘、棋子、提示信息及按钮"""

    def __init__(self, screen, game_core, interaction_handler):
        self.screen = screen  # pygame屏幕对象
        self.game_core = game_core  # 关联核心游戏模块
        self.interaction = interaction_handler  # 关联交互模块

        # 按钮区域定义 - 两行布局，靠右对齐
        # 按钮宽度和间距设置
        button_width = 100
        button_height = 40
        button_margin = 20
        
        # 计算右侧起始位置，确保按钮靠右对齐
        right_margin = MARGIN
        
        self.buttons = {
            # 第一行（靠下）：保存棋局、加载棋局、退出游戏
            "save": pygame.Rect(SCREEN_WIDTH - right_margin - button_width * 3 - button_margin * 2, SCREEN_HEIGHT - 70, button_width, button_height),
            "load": pygame.Rect(SCREEN_WIDTH - right_margin - button_width * 2 - button_margin, SCREEN_HEIGHT - 70, button_width, button_height),
            "exit": pygame.Rect(SCREEN_WIDTH - right_margin - button_width, SCREEN_HEIGHT - 70, button_width, button_height),
            # 第二行（靠上）：重新开始、悔棋、返回主页面
            "restart": pygame.Rect(SCREEN_WIDTH - right_margin - button_width * 3 - button_margin * 2, SCREEN_HEIGHT - 130, button_width, button_height),
            "undo": pygame.Rect(SCREEN_WIDTH - right_margin - button_width * 2 - button_margin, SCREEN_HEIGHT - 130, button_width, button_height),
            "back_to_menu": pygame.Rect(SCREEN_WIDTH - right_margin - button_width, SCREEN_HEIGHT - 130, button_width, button_height),
            # 模式选择按钮
            "human_vs_human": pygame.Rect(200, 200, 200, 50),
            "human_vs_ai": pygame.Rect(200, 300, 200, 50),
            # 难度选择按钮（两行排列）
            "easy": pygame.Rect(100, 200, 100, 50),
            "normal": pygame.Rect(230, 200, 100, 50),
            "medium": pygame.Rect(360, 200, 100, 50),
            "hard": pygame.Rect(100, 280, 100, 50),
            "expert": pygame.Rect(230, 280, 100, 50),
        }

        # 字体初始化
        self.font = pygame.font.SysFont("SimHei", 24)
        self.small_font = pygame.font.SysFont("SimHei", 18)
        self.large_font = pygame.font.SysFont("SimHei", 36)
        
        # 动画相关
        self.thinking_animation_frame = 0
        self.last_animation_time = time.time()

    def render_all(self, show_mode_selection=False, show_difficulty_selection=False):
        """整合所有绘制逻辑，刷新屏幕
        
        Args:
            show_mode_selection: 是否显示模式选择界面
            show_difficulty_selection: 是否显示难度选择界面
        """
        if show_mode_selection:
            self.draw_mode_selection()
        elif show_difficulty_selection:
            self.draw_difficulty_selection()
        else:
            self.draw_board()
            self.draw_stones()
            self.draw_ui_elements()
            if self.game_core.ai_thinking:
                self.draw_thinking_animation()
        
        pygame.display.flip()

    def draw_board(self):
        """优化版绘制棋盘背景、网格线，提高渲染效率"""
        # 填充背景色
        self.screen.fill(BACKGROUND_COLOR)

        # 优化：预计算减少重复计算
        start_pos = MARGIN
        end_pos = MARGIN + (BOARD_SIZE - 1) * CELL_SIZE

        # 绘制网格线
        for i in range(BOARD_SIZE):
            # 横线
            pygame.draw.line(self.screen, LINE_COLOR,
                             (start_pos, start_pos + i * CELL_SIZE),
                             (end_pos, start_pos + i * CELL_SIZE),
                             2)
            # 竖线
            pygame.draw.line(self.screen, LINE_COLOR,
                             (start_pos + i * CELL_SIZE, start_pos),
                             (start_pos + i * CELL_SIZE, end_pos),
                             2)

    def draw_stones(self):
        """优化版绘制黑白棋子，使用placed_stones列表提高效率并标记最后落子位置"""
        # 只绘制已放置的棋子，避免遍历整个棋盘
        for row, col, player in self.game_core.placed_stones:
            color = BLACK_STONE if player == 1 else WHITE_STONE
            pos = (MARGIN + col * CELL_SIZE, MARGIN + row * CELL_SIZE)
            # 绘制棋子
            pygame.draw.circle(self.screen, color, pos, CELL_SIZE // 2 - 2)
            # 绘制棋子边缘
            pygame.draw.circle(self.screen, LINE_COLOR, pos, CELL_SIZE // 2 - 2, 1)
            
            # 标记最后落子位置
            if self.game_core.last_move and row == self.game_core.last_move[0] and col == self.game_core.last_move[1]:
                pygame.draw.circle(self.screen, LAST_MOVE_MARKER_COLOR, pos, LAST_MOVE_MARKER_RADIUS)

    def draw_ui_elements(self):
        """绘制提示信息、操作按钮"""
        # 绘制状态提示
        text = self.font.render(self.game_core.message, True, TEXT_COLOR)
        self.screen.blit(text, (MARGIN, SCREEN_HEIGHT - 180))

        # 绘制键盘输入提示
        if self.interaction.waiting_for_input:
            input_text = self.font.render(f"输入: {self.interaction.keyboard_input}", True, TEXT_COLOR)
            self.screen.blit(input_text, (MARGIN, SCREEN_HEIGHT - 140))

        # 绘制文件名输入提示（当处于输入状态时）
        if self.interaction.waiting_for_filename:
            input_prompt = self.font.render("请输入保存的文件名（按回车确认，ESC取消）：", True, TEXT_COLOR)
            self.screen.blit(input_prompt, (MARGIN, SCREEN_HEIGHT - 180))

            input_text = self.font.render(f"文件名: {self.interaction.filename_input}", True, TEXT_COLOR)
            self.screen.blit(input_text, (MARGIN, SCREEN_HEIGHT - 150))

        # 绘制加载文件名输入提示
        if self.interaction.waiting_for_load_filename:
            input_prompt = self.font.render("请输入要加载的文件名（按回车确认，ESC取消）：", True, TEXT_COLOR)
            self.screen.blit(input_prompt, (MARGIN, SCREEN_HEIGHT - 180))

            input_text = self.font.render(f"文件名: {self.interaction.filename_input}", True, TEXT_COLOR)
            self.screen.blit(input_text, (MARGIN, SCREEN_HEIGHT - 150))

        # 绘制按钮
        for name, rect in self.buttons.items():
            # 只绘制游戏界面相关的按钮，不绘制模式选择按钮
            if name in ["restart", "save", "load", "exit", "undo", "back_to_menu"]:
                # 按钮背景色（根据鼠标悬停状态）
                if rect.collidepoint(pygame.mouse.get_pos()):
                    color = BUTTON_HOVER
                else:
                    color = BUTTON_COLOR

                # 绘制按钮矩形
                pygame.draw.rect(self.screen, color, rect, border_radius=5)

                # 按钮文字
                button_text = {
                    "restart": "重新开始",
                    "save": "保存棋局",
                    "load": "加载棋局",
                    "exit": "退出游戏",
                    "undo": "悔棋",
                    "back_to_menu": "返回主页面"
                }[name]

                text = self.small_font.render(button_text, True, BUTTON_TEXT_COLOR)
                text_rect = text.get_rect(center=rect.center)
                self.screen.blit(text, text_rect)
    
    def draw_mode_selection(self):
        """绘制游戏模式选择界面"""
        # 填充背景
        self.screen.fill(BACKGROUND_COLOR)
        
        # 绘制标题
        title = self.large_font.render("请选择游戏模式", True, TEXT_COLOR)
        title_rect = title.get_rect(center=(SCREEN_WIDTH // 2, 100))
        self.screen.blit(title, title_rect)
        
        # 绘制按钮
        for name in ["human_vs_human", "human_vs_ai"]:
            rect = self.buttons[name]
            # 按钮背景色（根据鼠标悬停状态）
            if rect.collidepoint(pygame.mouse.get_pos()):
                color = BUTTON_HOVER
            else:
                color = BUTTON_COLOR
            
            # 绘制按钮矩形
            pygame.draw.rect(self.screen, color, rect, border_radius=5)
            
            # 按钮文字
            button_text = {
                "human_vs_human": "双人对战",
                "human_vs_ai": "人机对战"
            }[name]
            
            text = self.font.render(button_text, True, BUTTON_TEXT_COLOR)
            text_rect = text.get_rect(center=rect.center)
            self.screen.blit(text, text_rect)
    
    def draw_difficulty_selection(self):
        """绘制AI难度选择界面"""
        # 填充背景
        self.screen.fill(BACKGROUND_COLOR)
        
        # 绘制标题
        title = self.large_font.render("请选择AI难度", True, TEXT_COLOR)
        title_rect = title.get_rect(center=(SCREEN_WIDTH // 2, 100))
        self.screen.blit(title, title_rect)
        
        # 绘制提示
        hint = self.font.render("简单 - 普通 - 中等 - 困难 - 专家", True, TEXT_COLOR)
        hint_rect = hint.get_rect(center=(SCREEN_WIDTH // 2, 150))
        self.screen.blit(hint, hint_rect)
        
        # 绘制按钮
        for name in ["easy", "normal", "medium", "hard", "expert"]:
            rect = self.buttons[name]
            # 按钮背景色（根据鼠标悬停状态）
            if rect.collidepoint(pygame.mouse.get_pos()):
                color = BUTTON_HOVER
            else:
                color = BUTTON_COLOR
            
            # 绘制按钮矩形
            pygame.draw.rect(self.screen, color, rect, border_radius=5)
            
            # 按钮文字
            button_text = {
                "easy": "简单",
                "normal": "普通",
                "medium": "中等",
                "hard": "困难",
                "expert": "专家"
            }[name]
            
            text = self.font.render(button_text, True, BUTTON_TEXT_COLOR)
            text_rect = text.get_rect(center=rect.center)
            self.screen.blit(text, text_rect)
    
    def draw_thinking_animation(self):
        """绘制AI思考动画"""
        # 更新动画帧
        current_time = time.time()
        if current_time - self.last_animation_time > 0.5:  # 每0.5秒更新一帧
            self.thinking_animation_frame = (self.thinking_animation_frame + 1) % 4
            self.last_animation_time = current_time
        
        # 绘制思考图标（闪烁的省略号）
        dots = "." * (self.thinking_animation_frame + 1)
        text = self.font.render(f"AI思考中{dots}", True, TEXT_COLOR)
        self.screen.blit(text, (MARGIN, SCREEN_HEIGHT - 140))


class DataManager:
    """数据管理模块：负责棋局的保存与加载"""

    def __init__(self):
        self.default_dir = "."  # 棋局保存默认目录（当前工作目录）
        self.file_ext = ".txt"  # 棋局文件扩展名
        # 文件选择对话框相关配置
        self.file_filters = [("五子棋棋局文件", f"*{self.file_ext}"), ("所有文件", "*.*")]
        self.has_tkinter = self.check_tkinter()

    def check_tkinter(self):
        """检查是否有tkinter模块"""
        try:
            import tkinter as tk
            return True
        except ImportError:
            return False

    def save_game(self, game_core, custom_filename=None):
        """将当前棋局状态写入文本文件"""
        # 处理文件名
        if custom_filename:
            # 去除非法字符
            valid_chars = "-_.() abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
            filename = ''.join(c for c in custom_filename if c in valid_chars)
            # 添加扩展名
            if not filename.endswith(self.file_ext):
                filename += self.file_ext
        else:
            # 生成时间戳文件名
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"gobang_{timestamp}{self.file_ext}"

        # 检查文件是否存在，处理重名
        full_path = os.path.join(self.default_dir, filename)
        counter = 1
        while os.path.exists(full_path):
            # 询问用户是否覆盖
            if self.ask_override(filename):
                break
            # 不覆盖则添加序号
            name, ext = os.path.splitext(filename)
            filename = f"{name}({counter}){ext}"
            full_path = os.path.join(self.default_dir, filename)
            counter += 1

        # 写入文件
        try:
            with open(full_path, 'w', encoding='utf-8') as f:
                # 保存当前玩家
                f.write(f"当前玩家: {'黑方' if game_core.current_player == 1 else '白方'}\n")
                # 保存游戏状态
                f.write(f"游戏状态: {'已结束' if game_core.game_over else '进行中'}\n")
                if game_core.game_over:
                    if game_core.winner != 0:
                        f.write(f"获胜方: {'黑方' if game_core.winner == 1 else '白方'}\n")
                    else:
                        f.write("获胜方: 无（平局）\n")
                # 保存棋盘状态
                f.write("棋盘状态:\n")
                for row in game_core.board:
                    f.write(' '.join(map(str, row)) + '\n')

            return filename  # 返回实际保存的文件名
        except Exception as e:
            raise Exception(f"保存失败: {str(e)}")

    def load_game(self):
        """从文件读取棋局状态并恢复游戏（使用文件选择对话框）"""
        if not self.has_tkinter:
            # 没有tkinter，提示用户使用命令行输入文件名
            raise Exception("缺少tkinter模块，无法打开文件选择对话框。请使用键盘输入文件名进行加载。")

        # 弹出文件选择对话框
        try:
            import tkinter as tk
            from tkinter import filedialog
        except ImportError:
            raise Exception("加载失败: 缺少tkinter模块")

        root = tk.Tk()
        root.withdraw()  # 隐藏主窗口
        file_path = filedialog.askopenfilename(
            title="选择棋局文件",
            initialdir=self.default_dir,
            filetypes=self.file_filters
        )
        root.destroy()

        if not file_path:  # 用户取消选择
            return None

        return self.parse_game_file(file_path)

    def load_game_by_filename(self, filename):
        """通过文件名加载棋局（用于没有tkinter的情况）"""
        # 处理文件名
        if not filename.endswith(self.file_ext):
            filename += self.file_ext

        full_path = os.path.join(self.default_dir, filename)
        return self.parse_game_file(full_path)

    def parse_game_file(self, file_path):
        """解析游戏文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                lines = [line.strip() for line in f if line.strip()]

            if not lines:
                raise Exception("文件为空，无法加载")

            # 解析当前玩家
            current_player_line = next((line for line in lines if line.startswith("当前玩家:")), None)
            if not current_player_line:
                raise Exception("文件格式错误，缺少当前玩家信息")

            current_player_str = current_player_line.split(":", 1)[1].strip()
            if current_player_str == "黑方":
                current_player = 1
            elif current_player_str == "白方":
                current_player = 2
            else:
                raise Exception("无效的玩家信息")

            # 解析游戏状态
            game_status_line = next((line for line in lines if line.startswith("游戏状态:")), None)
            game_over = False
            if game_status_line:
                game_status = game_status_line.split(":", 1)[1].strip()
                game_over = (game_status == "已结束")

            # 解析获胜方
            winner_line = next((line for line in lines if line.startswith("获胜方:")), None)
            winner = 0
            if winner_line:
                winner_str = winner_line.split(":", 1)[1].strip()
                if winner_str == "黑方":
                    winner = 1
                elif winner_str == "白方":
                    winner = 2

            # 解析棋盘状态
            board_start_idx = None
            for i, line in enumerate(lines):
                if line.startswith("棋盘状态:"):
                    board_start_idx = i + 1
                    break

            if board_start_idx is None or board_start_idx >= len(lines):
                raise Exception("文件格式错误，缺少棋盘状态")

            board_lines = lines[board_start_idx:board_start_idx + BOARD_SIZE]
            if len(board_lines) != BOARD_SIZE:
                raise Exception(f"棋盘数据不完整，需要{BOARD_SIZE}行")

            # 解析棋盘数据
            board = []
            for i, line in enumerate(board_lines):
                cells = list(map(int, line.split()))
                if len(cells) != BOARD_SIZE:
                    raise Exception(f"第{i + 1}行棋盘数据格式错误")

                for cell in cells:
                    if cell not in (0, 1, 2):
                        raise Exception(f"棋盘数据包含非法值: {cell}")
                board.append(cells)

            return {
                'board': board,
                'current_player': current_player,
                'game_over': game_over,
                'winner': winner,
                'filename': os.path.basename(file_path)
            }

        except FileNotFoundError:
            raise Exception(f"文件不存在: {file_path}")
        except Exception as e:
            raise Exception(f"加载失败: {str(e)}")

    def ask_override(self, filename):
        """弹窗询问是否覆盖已有文件"""
        if not self.has_tkinter:
            # 没有tkinter，默认不覆盖
            return False

        try:
            import tkinter as tk
            from tkinter import messagebox
        except ImportError:
            return False

        root = tk.Tk()
        root.withdraw()
        result = messagebox.askyesno(
            "文件已存在",
            f"文件 '{filename}' 已存在，是否覆盖？"
        )
        root.destroy()
        return result

    def list_save_files(self):
        """列出当前目录下的所有保存文件"""
        save_files = []
        for file in os.listdir(self.default_dir):
            if file.endswith(self.file_ext):
                # 检查文件名是否包含游戏相关的关键词
                if any(keyword in file.lower() for keyword in ['gobang', 'game', 'save']):
                    save_files.append(file)
        return save_files


class MainController:
    """主控制模块：协调各模块工作，处理事件循环"""

    def __init__(self):
        # 初始化pygame
        pygame.init()
        pygame.display.set_caption("五子棋 - AI对战")

        # 创建屏幕
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))

        # 初始化各模块
        self.game_core = GameCore()
        self.interaction = InteractionHandler(self.game_core)
        self.data_manager = DataManager()
        self.renderer = UIRenderer(self.screen, self.game_core, self.interaction)

        # 设置模块间引用
        self.interaction.set_data_manager(self.data_manager)

        self.running = True  # 游戏运行标识
        self.show_mode_selection = True  # 是否显示模式选择界面
        self.show_difficulty_selection = False  # 是否显示难度选择界面
        self.last_ai_move_time = 0  # 记录上次AI落子时间

    def run(self):
        """游戏主循环"""
        clock = pygame.time.Clock()

        while self.running:
            # 事件处理
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    self.running = False
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    if event.button == 1:  # 左键点击
                        self.handle_mouse_click(event.pos)
                elif event.type == pygame.KEYDOWN:
                    self.interaction.handle_keyboard_event(event)
                elif event.type == pygame.USEREVENT:
                    # 处理AI落子事件
                    if hasattr(event, 'move') and not self.game_core.game_over:
                        row, col = event.move
                        # 执行AI落子
                        self.game_core.make_move(row, col)

            # 如果是人机对战且当前是AI回合且AI不在思考中
            if (not self.show_mode_selection and 
                not self.show_difficulty_selection and 
                self.game_core.game_mode == MODE_HUMAN_VS_AI and 
                self.game_core.current_player == self.game_core.ai_player.player and 
                not self.game_core.game_over and 
                not self.game_core.ai_thinking):
                # 添加短暂延迟，让用户看清状态变化
                current_time = time.time()
                if current_time - self.last_ai_move_time > 0.5:
                    self.game_core.ai_move()
                    self.last_ai_move_time = current_time

            # 渲染界面
            self.renderer.render_all(self.show_mode_selection, self.show_difficulty_selection)
            clock.tick(30)  # 控制帧率

        # 退出游戏
        pygame.quit()
        sys.exit()

    def handle_mouse_click(self, pos):
        """处理鼠标点击（区分棋盘与按钮）"""
        # 如果在模式选择界面
        if self.show_mode_selection:
            for name in ["human_vs_human", "human_vs_ai"]:
                if self.renderer.buttons[name].collidepoint(pos):
                    self.handle_button_click(name)
                    return
        # 如果在难度选择界面
        elif self.show_difficulty_selection:
            for name in ["easy", "normal", "medium", "hard", "expert"]:
                if self.renderer.buttons[name].collidepoint(pos):
                    self.handle_button_click(name)
                    return
        # 否则在游戏界面
        else:
            # 检查是否点击按钮
            for name in ["restart", "save", "load", "exit", "undo", "back_to_menu"]:
                if self.renderer.buttons[name].collidepoint(pos):
                    self.handle_button_click(name)
                    return

            # 否则视为棋盘点击
            self.interaction.handle_mouse_click(pos)

    def handle_button_click(self, name):
        """处理按钮点击事件"""
        if name == "restart":
            self.game_core.reset_game()
        elif name == "save":
            # 进入文件名输入状态
            self.interaction.waiting_for_filename = True
            self.interaction.filename_input = ""
            self.game_core.message = "请输入保存的文件名（按回车确认，ESC取消）"
        elif name == "load":
            if self.data_manager.has_tkinter:
                # 调用加载方法并处理结果
                try:
                    result = self.data_manager.load_game()
                    if result:
                        self.game_core.board = result['board']
                        self.game_core.current_player = result['current_player']
                        self.game_core.game_over = result['game_over']
                        self.game_core.winner = result['winner']
                        
                        # 根据棋盘状态重建其他属性
                        self.game_core.move_history = []
                        self.game_core.placed_stones = []
                        self.game_core.stone_count = 0
                        self.game_core.last_move = None
                        
                        # 遍历棋盘重建历史记录和其他属性
                        for row_idx, row in enumerate(self.game_core.board):
                            for col_idx, cell in enumerate(row):
                                if cell != 0:  # 非空位置
                                    self.game_core.move_history.append((row_idx, col_idx, cell))
                                    self.game_core.placed_stones.append((row_idx, col_idx, cell))
                                    self.game_core.stone_count += 1
                                    # 如果是最后一个落子，更新last_move
                                    if self.game_core.stone_count == len(self.game_core.move_history):
                                        self.game_core.last_move = (row_idx, col_idx)

                        if self.game_core.game_over:
                            if self.game_core.winner != 0:
                                self.game_core.message = f"已加载棋局: {result['filename']}，{'黑方' if self.game_core.winner == 1 else '白方'}获胜"
                            else:
                                self.game_core.message = f"已加载棋局: {result['filename']}，平局"
                        else:
                            if self.game_core.game_mode == MODE_HUMAN_VS_AI:
                                if self.game_core.current_player == 1:
                                    self.game_core.message = f"已加载棋局: {result['filename']}，黑方回合（您）"
                                else:
                                    self.game_core.message = f"已加载棋局: {result['filename']}，AI思考中..."
                            else:
                                self.game_core.message = f"已加载棋局: {result['filename']}，当前为{'黑方' if self.game_core.current_player == 1 else '白方'}回合"
                    else:
                        self.game_core.message = "取消加载"
                except Exception as e:
                    self.game_core.message = str(e)
            else:
                # 没有tkinter，使用键盘输入文件名
                self.interaction.waiting_for_load_filename = True
                self.interaction.filename_input = ""
                # 列出当前目录下的保存文件
                save_files = self.data_manager.list_save_files()
                if save_files:
                    files_str = "、".join(save_files[:5])  # 只显示前5个文件
                    if len(save_files) > 5:
                        files_str += "..."
                    self.game_core.message = f"请输入要加载的文件名（按回车确认，ESC取消）。当前目录下的棋局文件：{files_str}"
                else:
                    self.game_core.message = "请输入要加载的文件名（按回车确认，ESC取消）。当前目录下没有找到棋局文件。"
        elif name == "exit":
            self.running = False
        elif name == "back_to_menu":
            # 返回主页面，重置游戏并显示模式选择界面
            self.game_core.reset_game()
            self.show_mode_selection = True
        elif name == "undo":
            # 执行悔棋
            if self.game_core.can_undo():
                self.game_core.undo_move()
                self.game_core.message = "已悔棋"
            else:
                self.game_core.message = "没有可悔的步骤"
        # 游戏模式选择按钮
        elif name == "human_vs_human":
            self.game_core.set_game_mode(MODE_HUMAN_VS_HUMAN)
            self.show_mode_selection = False
            self.show_difficulty_selection = False
        elif name == "human_vs_ai":
            self.show_mode_selection = False
            self.show_difficulty_selection = True
        # 难度选择按钮
        elif name in ["easy", "normal", "medium", "hard", "expert"]:
            self.game_core.set_game_mode(MODE_HUMAN_VS_AI, ai_difficulty=name)
            self.show_mode_selection = False
            self.show_difficulty_selection = False


if __name__ == "__main__":
    # 启动游戏
    game = MainController()
    game.run()